import os
from langchain_community.document_loaders import DirectoryLoader, TextLoader

"""加载目录中的文档时，如果DirectoryLoader函数不给参数loader_cls赋值，执行时会报zipfile.BadZipFile: File is not a zip file
网上找的如下所示的解决办法也不生效
import nltk
# Set NLTK data path
nltk.data.path.append("D:/nltk_data")
# Ensure NLTK data is available
required_nltk_data = [
    ('tokenizers/punkt', 'punkt'),
    ('taggers/averaged_perceptron_tagger', 'averaged_perceptron_tagger')
]

for path, package in required_nltk_data:
    try:
        nltk.data.find(path)
    except LookupError:
        print(f"Downloading NLTK data: {package}")
        nltk.download(package, download_dir="D:/nltk_data")
"""

# 获取当前脚本文件所在的目录
script_dir = os.path.dirname(__file__)
print(f"获取当前脚本文件所在的目录：{script_dir}")
data_dir = "D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空"
# loader = DirectoryLoader(data_dir, glob="*.txt", loader_cls=TextLoader, loader_kwargs={"encoding": "utf-8"})
# 设置show_progress=True可显示加载进度；设置silent_errors=True跳过无法加载的文件
loader = DirectoryLoader(data_dir, glob="**/*.md",
                         loader_cls=TextLoader, loader_kwargs={"encoding": "utf-8"},
                         use_multithreading=True, show_progress=True, silent_errors=True)
docs = loader.load()
print(f"文档数：{len(docs)}")  # 输出文档总数
print(docs[0])  # 输出第一个文档